#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
增强版账单分析程序
支持支付宝CSV和微信Excel格式的流水记录分析
"""

import pandas as pd
import os
import re
from datetime import datetime, timedelta
from typing import Dict, List, Tuple, Optional
from tabulate import tabulate
import warnings
warnings.filterwarnings('ignore')


class EnhancedBillAnalyzer:
    """增强版账单分析器"""
    
    def __init__(self, bill_folder: str = "zhangdan"):
        self.bill_folder = bill_folder
        self.user_nicknames = {}  # 存储用户昵称
        self.bills_data = []  # 存储处理后的账单数据
        self.field_mapping = self._init_field_mapping()
        self.category_rules = self._init_category_rules()
        
    def _init_field_mapping(self) -> Dict[str, Dict[str, str]]:
        """初始化字段映射关系"""
        return {
            'alipay': {
                '交易时间': 'datetime',
                '交易对方': 'counterparty', 
                '商品说明': 'description',
                '收/支': 'type',
                '金额': 'amount',
                '收/付款方式': 'payment_method',
                '交易状态': 'status',
                '交易分类': 'category',
                '对方账号': 'counterparty_account',
                '备注': 'remark'
            },
            'wechat': {
                '交易时间': 'datetime',
                '交易类型': 'type',
                '交易对方': 'counterparty',
                '商品': 'description', 
                '收/支': 'income_expense',
                '金额(元)': 'amount',
                '支付方式': 'payment_method',
                '当前状态': 'status',
                '交易单号': 'transaction_id',
                '商户单号': 'merchant_id',
                '备注': 'remark'
            }
        }
    
    def _init_category_rules(self) -> Dict[str, Dict[str, List[str]]]:
        """初始化分类规则"""
        return {
            '餐饮': {
                'keywords': ['餐饮', '外卖', '餐厅', '咖啡', '奶茶', '小吃', '美团', '饿了么', '星巴克', '麦当劳', '肯德基', '蜜雪冰城'],
                'subcategories': {
                    '外卖': ['外卖', '配送', '美团', '饿了么'],
                    '饮品': ['咖啡', '奶茶', '星巴克', '蜜雪冰城'],
                    '堂食': ['餐厅', '饭店', '餐馆']
                }
            },
            '交通': {
                'keywords': ['交通', '打车', '地铁', '公交', '加油', '停车', '滴滴', '高德', '地铁', '公交', '加油站'],
                'subcategories': {
                    '网约车': ['打车', '网约车', '滴滴', '高德'],
                    '公共交通': ['地铁', '公交'],
                    '汽车相关': ['加油', '停车', '加油站']
                }
            },
            '购物': {
                'keywords': ['购物', '商品', '购买', '订单', '淘宝', '京东', '拼多多', '天猫', '天猫优选'],
                'subcategories': {
                    '服装': ['服装', '衣服', '鞋', '服饰'],
                    '食品': ['食品', '零食', '水果', '生鲜'],
                    '数码产品': ['数码', '电子', '手机', '电脑', '手机壳'],
                    '其他商品': ['商品', '购买', '订单']
                }
            },
            '生活服务': {
                'keywords': ['水电', '燃气', '物业', '维修', '清洁', '电费', '水费', '物业费'],
                'subcategories': {
                    '水电费': ['电费', '水费', '燃气费'],
                    '物业费': ['物业', '物业费'],
                    '维修服务': ['维修', '清洁', '家政']
                }
            },
            '娱乐': {
                'keywords': ['娱乐', '游戏', '电影', '音乐', '视频', '腾讯', '爱奇艺', '优酷', '网易'],
                'subcategories': {
                    '游戏': ['游戏', '腾讯游戏', '网易游戏'],
                    '视频': ['电影', '视频', '爱奇艺', '优酷'],
                    '音乐': ['音乐', '网易云音乐']
                }
            },
            '医疗': {
                'keywords': ['医疗', '医院', '药店', '体检', '药品'],
                'subcategories': {
                    '医院': ['医院', '体检'],
                    '药店': ['药店', '药品']
                }
            },
            '教育': {
                'keywords': ['教育', '培训', '课程', '学习', '教育咨询'],
                'subcategories': {
                    '培训': ['培训', '课程', '教育咨询'],
                    '学习': ['学习', '教育']
                }
            },
            '投资理财': {
                'keywords': ['余额宝', '基金', '理财', '投资', '蚂蚁财富', '兴银货币'],
                'subcategories': {
                    '余额宝': ['余额宝', '兴银货币'],
                    '基金': ['基金', '蚂蚁财富'],
                    '其他理财': ['理财', '投资']
                }
            },
            '转账': {
                'keywords': ['转账', '红包', '收钱码', '充值'],
                'subcategories': {
                    '个人转账': ['转账', '个人'],
                    '红包': ['红包', '现金红包'],
                    '收付款': ['收钱码', '充值']
                }
            }
        }
    
    def load_bills(self) -> None:
        """加载账单文件"""
        if not os.path.exists(self.bill_folder):
            print(f"错误：账单文件夹 '{self.bill_folder}' 不存在")
            return
            
        files = os.listdir(self.bill_folder)
        alipay_files = [f for f in files if f.endswith('.csv')]
        wechat_files = [f for f in files if f.endswith('.xlsx')]
        
        print(f"发现 {len(alipay_files)} 个支付宝文件，{len(wechat_files)} 个微信文件")
        
        # 处理支付宝文件
        for file in alipay_files:
            self._process_alipay_file(file)
            
        # 处理微信文件
        for file in wechat_files:
            self._process_wechat_file(file)
    
    def _process_alipay_file(self, filename: str) -> None:
        """处理支付宝CSV文件"""
        filepath = os.path.join(self.bill_folder, filename)
        print(f"正在处理支付宝文件: {filename}")
        
        try:
            # 先读取文件内容找到数据开始行
            with open(filepath, 'r', encoding='gbk', errors='ignore') as f:
                lines = f.readlines()
            
            # 查找包含"交易时间"的行
            data_start_row = None
            for i, line in enumerate(lines):
                if '交易时间' in line and '交易分类' in line:
                    data_start_row = i
                    break
            
            if data_start_row is None:
                print(f"无法找到数据开始行，跳过文件 {filename}")
                return
            
            print(f"数据从第 {data_start_row + 1} 行开始")
            
            # 提取用户信息和时间范围
            self._extract_file_info(filepath, 'alipay')
            
            # 使用pandas读取数据，跳过前面的行
            df = pd.read_csv(filepath, encoding='gbk', skiprows=data_start_row)
            
            # 处理数据
            processed_data = self._process_alipay_data(df)
            self.bills_data.extend(processed_data)
            
        except Exception as e:
            print(f"处理支付宝文件 {filename} 时出错: {e}")
    
    def _process_wechat_file(self, filename: str) -> None:
        """处理微信Excel文件"""
        filepath = os.path.join(self.bill_folder, filename)
        print(f"正在处理微信文件: {filename}")
        
        try:
            # 读取Excel文件
            df = pd.read_excel(filepath)
            
            # 查找数据开始行
            data_start_row = self._find_data_start_row(df)
            if data_start_row is None:
                print(f"无法找到数据开始行，跳过文件 {filename}")
                return
                
            # 重新读取数据
            df = pd.read_excel(filepath, skiprows=data_start_row)
            
            # 提取用户信息和时间范围
            self._extract_file_info(filepath, 'wechat')
            
            # 处理数据
            processed_data = self._process_wechat_data(df)
            self.bills_data.extend(processed_data)
            
        except Exception as e:
            print(f"处理微信文件 {filename} 时出错: {e}")
    
    def _find_data_start_row(self, df: pd.DataFrame) -> Optional[int]:
        """查找数据开始行"""
        for i, row in df.iterrows():
            # 查找包含"交易时间"或"时间"的行
            row_str = ' '.join([str(cell) for cell in row.values if pd.notna(cell)])
            if '交易时间' in row_str or '时间' in row_str:
                return i
        return None
    
    def _extract_file_info(self, filepath: str, platform: str) -> None:
        """提取文件信息（用户昵称、时间范围等）"""
        try:
            with open(filepath, 'r', encoding='gbk', errors='ignore') as f:
                content = f.read()
                
            # 提取时间范围
            time_pattern = r'(\d{4}-\d{2}-\d{2})'
            times = re.findall(time_pattern, content)
            if len(times) >= 2:
                start_time = times[0]
                end_time = times[1]
                print(f"时间范围: {start_time} 到 {end_time}")
            
            # 提取用户信息
            user_key = f"{platform}_user"
            if user_key not in self.user_nicknames:
                # 尝试从文件中提取用户名
                name_match = re.search(r'姓名：(.+)', content)
                if name_match:
                    self.user_nicknames[user_key] = name_match.group(1).strip()
                else:
                    self.user_nicknames[user_key] = f"用户_{len(self.user_nicknames) + 1}"
                
        except Exception as e:
            print(f"提取文件信息时出错: {e}")
    
    def _process_alipay_data(self, df: pd.DataFrame) -> List[Dict]:
        """处理支付宝数据"""
        processed_data = []
        
        for _, row in df.iterrows():
            if pd.isna(row.get('交易时间')):
                continue
                
            try:
                # 解析时间
                datetime_str = str(row['交易时间'])
                if ' ' in datetime_str:
                    dt = datetime.strptime(datetime_str, '%Y-%m-%d %H:%M:%S')
                else:
                    dt = datetime.strptime(datetime_str, '%Y-%m-%d')
                
                # 处理金额
                amount_str = str(row.get('金额', '0')).replace(',', '')
                amount = float(amount_str) if amount_str.replace('.', '').isdigit() else 0.0
                
                # 确定收支类型
                income_expense = str(row.get('收/支', ''))
                if '收入' in income_expense or '收款' in income_expense:
                    bill_type = 'income'
                elif '支出' in income_expense or '付款' in income_expense:
                    bill_type = 'expense'
                else:
                    bill_type = 'unknown'
                
                bill_data = {
                    'platform': 'alipay',
                    'datetime': dt,
                    'counterparty': str(row.get('交易对方', '')),
                    'description': str(row.get('商品说明', '')),
                    'type': bill_type,
                    'amount': amount,
                    'payment_method': str(row.get('收/付款方式', '')),
                    'status': str(row.get('交易状态', '')),
                    'category': str(row.get('交易分类', '')),
                    'remark': str(row.get('备注', '')),
                    'raw_data': row.to_dict()
                }
                
                processed_data.append(bill_data)
                
            except Exception as e:
                print(f"处理支付宝数据行时出错: {e}")
                continue
                
        return processed_data
    
    def _process_wechat_data(self, df: pd.DataFrame) -> List[Dict]:
        """处理微信数据"""
        processed_data = []
        
        for _, row in df.iterrows():
            if pd.isna(row.get('交易时间')):
                continue
                
            try:
                # 解析时间
                datetime_str = str(row['交易时间'])
                if ' ' in datetime_str:
                    dt = datetime.strptime(datetime_str, '%Y-%m-%d %H:%M:%S')
                else:
                    dt = datetime.strptime(datetime_str, '%Y-%m-%d')
                
                # 处理金额
                amount_str = str(row.get('金额(元)', '0')).replace(',', '')
                amount = float(amount_str) if amount_str.replace('.', '').isdigit() else 0.0
                
                # 确定收支类型
                income_expense = str(row.get('收/支', ''))
                if '收入' in income_expense or '收款' in income_expense:
                    bill_type = 'income'
                elif '支出' in income_expense or '付款' in income_expense:
                    bill_type = 'expense'
                else:
                    bill_type = 'unknown'
                
                bill_data = {
                    'platform': 'wechat',
                    'datetime': dt,
                    'counterparty': str(row.get('交易对方', '')),
                    'description': str(row.get('商品', '')),
                    'type': bill_type,
                    'amount': amount,
                    'payment_method': str(row.get('支付方式', '')),
                    'status': str(row.get('当前状态', '')),
                    'remark': str(row.get('备注', '')),
                    'raw_data': row.to_dict()
                }
                
                processed_data.append(bill_data)
                
            except Exception as e:
                print(f"处理微信数据行时出错: {e}")
                continue
                
        return processed_data
    
    def classify_bills(self) -> None:
        """对账单进行分类"""
        for bill in self.bills_data:
            # 第一级分类
            bill['category_level1'] = self._classify_level1(bill)
            
            # 第二级分类
            bill['category_level2'] = self._classify_level2(bill)
            
            # 第三级分类
            bill['category_level3'] = self._classify_level3(bill)
    
    def _classify_level1(self, bill: Dict) -> str:
        """第一级分类"""
        description = bill['description'].lower()
        counterparty = bill['counterparty'].lower()
        text_to_check = f"{description} {counterparty}"
        
        for category, rules in self.category_rules.items():
            if any(keyword in text_to_check for keyword in rules['keywords']):
                return category
        
        return '其他'
    
    def _classify_level2(self, bill: Dict) -> str:
        """第二级分类"""
        level1 = bill['category_level1']
        description = bill['description'].lower()
        counterparty = bill['counterparty'].lower()
        text_to_check = f"{description} {counterparty}"
        
        if level1 in self.category_rules:
            subcategories = self.category_rules[level1]['subcategories']
            for subcategory, keywords in subcategories.items():
                if any(keyword in text_to_check for keyword in keywords):
                    return subcategory
        
        return '未分类'
    
    def _classify_level3(self, bill: Dict) -> str:
        """第三级分类"""
        level2 = bill['category_level2']
        counterparty = bill['counterparty']
        
        # 根据交易对方进行第三级分类
        if level2 == '外卖':
            if '美团' in counterparty:
                return '美团外卖'
            elif '饿了么' in counterparty:
                return '饿了么'
            else:
                return '其他外卖'
        
        elif level2 == '网约车':
            if '滴滴' in counterparty:
                return '滴滴出行'
            elif '高德' in counterparty:
                return '高德打车'
            else:
                return '其他网约车'
        
        elif level2 == '服装':
            if '淘宝' in counterparty:
                return '淘宝服装'
            elif '京东' in counterparty:
                return '京东服装'
            else:
                return '其他服装'
        
        elif level2 == '余额宝':
            if '兴银货币' in counterparty:
                return '兴银货币A'
            else:
                return '余额宝其他'
        
        else:
            return '未分类'
    
    def get_recent_bills(self, days: int = 20) -> List[Dict]:
        """获取最近N天的账单"""
        cutoff_date = datetime.now() - timedelta(days=days)
        recent_bills = [bill for bill in self.bills_data if bill['datetime'] >= cutoff_date]
        return sorted(recent_bills, key=lambda x: x['datetime'], reverse=True)
    
    def get_bills_by_month(self, year: int, month: int) -> List[Dict]:
        """获取指定年月的账单"""
        month_bills = [bill for bill in self.bills_data 
                      if bill['datetime'].year == year and bill['datetime'].month == month]
        return sorted(month_bills, key=lambda x: x['datetime'], reverse=True)
    
    def get_category_summary(self, bills: List[Dict]) -> Dict:
        """获取分类汇总"""
        summary = {}
        for bill in bills:
            level1 = bill['category_level1']
            if level1 not in summary:
                summary[level1] = {'income': 0, 'expense': 0, 'count': 0}
            
            if bill['type'] == 'income':
                summary[level1]['income'] += bill['amount']
            elif bill['type'] == 'expense':
                summary[level1]['expense'] += bill['amount']
            
            summary[level1]['count'] += 1
        
        return summary
    
    def display_bills(self, days: int = 20) -> None:
        """显示账单数据"""
        recent_bills = self.get_recent_bills(days)
        
        if not recent_bills:
            print(f"没有找到最近{days}天的账单数据")
            return
        
        print(f"\n最近 {days} 天的账单数据 (共 {len(recent_bills)} 条记录):")
        print("=" * 120)
        
        # 准备表格数据
        table_data = []
        for bill in recent_bills:
            table_data.append([
                bill['datetime'].strftime('%Y-%m-%d %H:%M'),
                '支付宝' if bill['platform'] == 'alipay' else '微信',
                bill['counterparty'][:15] + '...' if len(bill['counterparty']) > 15 else bill['counterparty'],
                bill['description'][:20] + '...' if len(bill['description']) > 20 else bill['description'],
                '收入' if bill['type'] == 'income' else '支出',
                f"{bill['amount']:.2f}",
                bill['category_level1'],
                bill['category_level2'],
                bill['category_level3']
            ])
        
        # 使用tabulate显示表格
        headers = ['时间', '平台', '交易对方', '商品描述', '类型', '金额', '一级分类', '二级分类', '三级分类']
        print(tabulate(table_data, headers=headers, tablefmt='grid'))
        
        # 统计信息
        total_income = sum(bill['amount'] for bill in recent_bills if bill['type'] == 'income')
        total_expense = sum(bill['amount'] for bill in recent_bills if bill['type'] == 'expense')
        
        print(f"\n统计信息:")
        print(f"总收入: {total_income:.2f} 元")
        print(f"总支出: {total_expense:.2f} 元")
        print(f"净收入: {total_income - total_expense:.2f} 元")
        
        # 分类汇总
        print(f"\n分类汇总:")
        category_summary = self.get_category_summary(recent_bills)
        for category, data in sorted(category_summary.items()):
            print(f"{category}: 收入 {data['income']:.2f} 元, 支出 {data['expense']:.2f} 元, 共 {data['count']} 笔")
    
    def display_monthly_summary(self, year: int, month: int) -> None:
        """显示月度汇总"""
        month_bills = self.get_bills_by_month(year, month)
        
        if not month_bills:
            print(f"没有找到 {year}年{month}月 的账单数据")
            return
        
        print(f"\n{year}年{month}月 账单汇总 (共 {len(month_bills)} 条记录):")
        print("=" * 80)
        
        # 统计信息
        total_income = sum(bill['amount'] for bill in month_bills if bill['type'] == 'income')
        total_expense = sum(bill['amount'] for bill in month_bills if bill['type'] == 'expense')
        
        print(f"总收入: {total_income:.2f} 元")
        print(f"总支出: {total_expense:.2f} 元")
        print(f"净收入: {total_income - total_expense:.2f} 元")
        
        # 分类汇总
        print(f"\n分类汇总:")
        category_summary = self.get_category_summary(month_bills)
        for category, data in sorted(category_summary.items()):
            print(f"{category}: 收入 {data['income']:.2f} 元, 支出 {data['expense']:.2f} 元, 共 {data['count']} 笔")
    
    def run(self) -> None:
        """运行分析程序"""
        print("开始账单分析...")
        
        # 加载账单
        self.load_bills()
        
        if not self.bills_data:
            print("没有找到有效的账单数据")
            return
        
        print(f"成功加载 {len(self.bills_data)} 条账单记录")
        
        # 分类账单
        self.classify_bills()
        
        # 显示结果
        self.display_bills(days=365)  # 显示所有数据
        
        # 显示月度汇总
        if self.bills_data:
            # 获取数据的时间范围
            dates = [bill['datetime'] for bill in self.bills_data]
            min_date = min(dates)
            max_date = max(dates)
            
            print(f"\n数据时间范围: {min_date.strftime('%Y-%m')} 到 {max_date.strftime('%Y-%m')}")
            
            # 显示每个月的汇总
            current_date = min_date.replace(day=1)
            while current_date <= max_date:
                self.display_monthly_summary(current_date.year, current_date.month)
                if current_date.month == 12:
                    current_date = current_date.replace(year=current_date.year + 1, month=1)
                else:
                    current_date = current_date.replace(month=current_date.month + 1)


def main():
    """主函数"""
    analyzer = EnhancedBillAnalyzer()
    analyzer.run()


if __name__ == "__main__":
    main()
